use super::*;
use riir::agent;
use riir::agent::tools::playground::{run_rustycli, VerifyResp, VerifyType};
use std::marker::PhantomData;
use tracing::info;

use serde::{Deserialize, Serialize};


#[derive(Serialize, Deserialize)]
pub struct RustAuditInfo {
    pub rust_expert: String,
}

pub const RUST_AUDIT: Prompt<RustAuditInfo> = Prompt("RustAuditInfo: [Rust Audit]", PhantomData);

#[derive(Serialize, Deserialize)]
pub struct AuditPlanInfo {
    pub task: String,
}

pub const RUST_AUDIT_PLAN: Prompt<AuditPlanInfo> = Prompt(
    r#"
You are a Rust expert assistant. You can extract Rust code from the user's input, interpret it, analyze it step by step, and then refactor it.

The requirements for audit the code are as follows:

- Please ensure the code is fully functional and doesn't omit any details.
- The strings in the code you generate do not need extra escape characters.
- Please review the code for potential concurrency deadlock issues.
- If it contains unsafe, please review whether the unsafe blocks satisfy the safety abstractions.
- Review the code for possible improvements.
- The strings in the code you generate do not need extra escape characters.
- Add comments and documentation to the code.
- The Rust code snippet after your refactoring must not have any escape characters, please!


If you discover deadlock in the code, you can fix it by following the methods below:
- Limit the loop to a finite number of iterations to prevent the infinite waiting condition.
- Release the lock more frequently to give other threads a chance to acquire the lock and proceed, reducing the likelihood of a deadlock.

After you have completed the review, please correct the code and output the complete code in the specified format.

Please ensure that your reply strictly adheres to the following JSON format:

```json
{{
	"plan": "<concise plan on how you will complete the task>",
	"code": "
	```rust
	<The Rust code after your audit, and the Rust code snippet without any escape characters Please!>
	```
	"
}}
```
"#,
    PhantomData,
);

#[derive(Debug)]
pub struct AuditAgent {
    attributes: AgentBase,
    bug_errors: Option<String>,
    bug_count: u8,
    retry_count: u8,
    llm: Arc<Mutex<Box<dyn LLM + 'static>>>,
}

impl AuditAgent {
    pub fn new(llm: Arc<Mutex<Box<dyn LLM + 'static>>>) -> Self {
        let attributes: AgentBase = AgentBase {
            objective: "Audit Rust Code".to_string(),
            position: "Audit Agent".to_string(),
            state: AgentState::Starting,
            memory: vec![],
        };

        Self {
            attributes,
            bug_errors: None,
            bug_count: 0,
            retry_count: 0,
            llm: llm.clone(),
        }
    }

    async fn call_initial_code(&mut self, factsheet: &mut FactSheet, output_path: PathBuf) {
        // let decode_code =
        //     decode_llm_output(factsheet.code.as_ref().unwrap_or(&"".to_string())).await;
        // Concatenate Instruction
        let msg_context: Vec<Message> = prompt::build_msg(
            format!("PLAN: {} ", RUST_AUDIT_PLAN.0),
            format!(
                "CODE: {} ",
                factsheet.code.as_ref().unwrap_or(&"".to_string())
            ),
        );

        let ai_response: String =
            ai_task_request(msg_context, &self.attributes.position, self.llm.clone()).await;

        save_code(output_path, "audit", &ai_response).await;
        factsheet.code = Some(ai_response);
    }

    async fn call_improved_code(&mut self, factsheet: &mut FactSheet, output_path: PathBuf) {
        // let decode_code =
        //     decode_llm_output(factsheet.code.as_ref().unwrap_or(&"".to_string())).await;
        let msg_context: Vec<Message> = prompt::build_msg(
            format!("PLAN: {} ", RUST_AUDIT_PLAN.0),
            format!(
                "CODE: {} ",
                factsheet.code.as_ref().unwrap_or(&"".to_string())
            ),
        );

        let ai_response: String =
            ai_task_request(msg_context, &self.attributes.position, self.llm.clone()).await;

        save_code(output_path, "audit", &ai_response).await;
        factsheet.code = Some(ai_response);
    }

    async fn call_fix_code_bugs(&mut self, factsheet: &mut FactSheet, output_path: PathBuf) {
        // let decode_code =
        //     decode_llm_output(factsheet.code.as_ref().unwrap_or(&"".to_string())).await;
        let msg_context: Vec<Message> = prompt::build_msg(
            format!("PLAN: {} ", RUST_FIXED_EXPERT.0),
            format!(
                "CODE: {:?} \n ERROR_BUGS: {:?} \n",
                factsheet.code.as_ref().unwrap_or(&"".to_string()),
                self.bug_errors
            ),
        );

        let ai_response: String =
            ai_task_request(msg_context, &self.attributes.position, self.llm.clone()).await;

        save_code(output_path, "audit", &ai_response).await;
        factsheet.code = Some(ai_response);
    }
}



#[async_trait]
impl AgentFns for AuditAgent {
    fn get_attributes_from_agent(&self) -> &AgentBase {
        &self.attributes
    }

    async fn execute(
        &mut self,
        factsheet: &mut FactSheet,
        output_path: PathBuf
    ) -> Result<(), Box<dyn std::error::Error>> {
        while self.attributes.state != AgentState::Done {
            match &self.attributes.state {
                AgentState::Starting => {
                    info!("Audit Agent Start ---");
                    self.call_initial_code(factsheet, output_path.clone()).await;
                    self.attributes.state = AgentState::Working;
                    continue;
                }

                AgentState::Working => {
                    info!("Audit Agent Working ---");
                    if self.bug_count == 0 {
                        self.call_improved_code(factsheet, output_path.clone()).await;
                    } else {
                        self.call_fix_code_bugs(factsheet, output_path.clone()).await;
                    }
                    self.attributes.state = AgentState::Debug;
                    continue;
                }

                AgentState::Debug => {
                    info!("Audit Agent Debug ---");

                    // --- 验证 LLM 响应代码
                    if let Some(e) = agent::capture_rust_code(&factsheet.code.clone().unwrap()) {
                        info!("EndANswer Chat: plain {:?}", e);

                        // do something async
                        let verft = run_rustycli(&e, VerifyType::BUILD).await;
                        info!("Verify Chat: {:?}", verft);
                        match verft {
                            VerifyResp::Success { stdout, stderr } => {
                                self.bug_count = 0;
                                info!("Code build successful...");
                                self.attributes.state = AgentState::Done;
                            }
                            VerifyResp::Fail(error_str) => {
                                // Update error stats
                                self.bug_count += 1;
                                self.retry_count += 1;
                                self.bug_errors = Some(error_str);
                                if self.retry_count > 10 {
                                    info!(
                                        "Code build failed but retry count is greater then 10... "
                                    );
                                    self.attributes.state = AgentState::Done;
                                } else {
                                    // Pass back for rework
                                    self.attributes.state = AgentState::Working;
                                    continue;
                                }
                            }
                        }
                    } else {
                        info!("Uncapture rust code!");
                        // Pass back for rework
                        self.attributes.state = AgentState::Working;
                        continue;
                    }
                }

                _ => {}
            }
        }
        Ok(())
    }
}
